#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版搜索工具
功能：在增强版知识库中进行高命中率搜索
"""

from pymilvus import connections, Collection
import jieba
import re
import sys


class EnhancedSearch:
    def __init__(self, collection_name="demo_enhanced_knowledge"):
        self.collection_name = collection_name
        self.collection = None
        
    def connect_to_milvus(self, host="localhost", port="19530"):
        """连接到Milvus"""
        try:
            connections.connect("default", host=host, port=port)
            self.collection = Collection(self.collection_name)
            self.collection.load()
            print(f"✅ 连接成功: {self.collection_name} (共{self.collection.num_entities}条记录)")
            return True
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return False
    
    def extract_search_terms(self, query):
        """提取搜索词"""
        # 产品编号
        products = re.findall(r'[A-Z]{1,3}\d{3,4}', query.upper())
        
        # 分词
        words = list(jieba.cut(query))
        keywords = [w for w in words if len(w) > 1 and w not in ['的', '是', '有', '和']]
        
        return products + keywords
    
    def multi_strategy_search(self, query, top_k=10):
        """多策略搜索"""
        search_terms = self.extract_search_terms(query)
        all_results = []
        
        print(f"🔍 搜索词: {search_terms}")
        
        # 策略1: 精确匹配 - 在原始文本中搜索
        try:
            conditions = [f'original_text like "%{term}%"' for term in search_terms[:3]]
            if conditions:
                expr = ' and '.join(conditions)
                results = self.collection.query(
                    expr=expr,
                    limit=top_k,
                    output_fields=['original_text', 'search_text', 'product_code', 
                                 'problem_category', 'solution', 'variant_type']
                )
                for result in results:
                    result['search_strategy'] = 'exact_match'
                    result['score'] = 10
                all_results.extend(results)
                print(f"📍 精确匹配: {len(results)} 条")
        except Exception as e:
            print(f"⚠️  精确匹配失败: {e}")
        
        # 策略2: 问题分类匹配
        try:
            for term in search_terms:
                if len(term) > 2:
                    expr = f'problem_category like "%{term}%"'
                    results = self.collection.query(
                        expr=expr,
                        limit=top_k//2,
                        output_fields=['original_text', 'search_text', 'product_code', 
                                     'problem_category', 'solution', 'variant_type']
                    )
                    for result in results:
                        result['search_strategy'] = 'category_match'
                        result['score'] = 8
                    all_results.extend(results)
            print(f"📂 分类匹配: {len([r for r in all_results if r.get('search_strategy') == 'category_match'])} 条")
        except Exception as e:
            print(f"⚠️  分类匹配失败: {e}")
        
        # 策略3: 解决方案匹配
        try:
            for term in search_terms:
                if len(term) > 2:
                    expr = f'solution like "%{term}%"'
                    results = self.collection.query(
                        expr=expr,
                        limit=top_k//2,
                        output_fields=['original_text', 'search_text', 'product_code', 
                                     'problem_category', 'solution', 'variant_type']
                    )
                    for result in results:
                        result['search_strategy'] = 'solution_match'
                        result['score'] = 7
                    all_results.extend(results)
            print(f"💡 解决方案匹配: {len([r for r in all_results if r.get('search_strategy') == 'solution_match'])} 条")
        except Exception as e:
            print(f"⚠️  解决方案匹配失败: {e}")
        
        # 策略4: 搜索文本匹配（包含同义词变体）
        try:
            conditions = [f'search_text like "%{term}%"' for term in search_terms[:3]]
            if conditions:
                expr = ' or '.join(conditions)
                results = self.collection.query(
                    expr=expr,
                    limit=top_k*2,
                    output_fields=['original_text', 'search_text', 'product_code', 
                                 'problem_category', 'solution', 'variant_type']
                )
                for result in results:
                    result['search_strategy'] = 'variant_match'
                    result['score'] = 5
                all_results.extend(results)
            print(f"🔄 变体匹配: {len([r for r in all_results if r.get('search_strategy') == 'variant_match'])} 条")
        except Exception as e:
            print(f"⚠️  变体匹配失败: {e}")
        
        # 去重和排序
        unique_results = []
        seen = set()
        
        for result in all_results:
            key = result['original_text']
            if key not in seen:
                seen.add(key)
                unique_results.append(result)
        
        # 按分数和策略排序
        strategy_priority = {'exact_match': 4, 'category_match': 3, 'solution_match': 2, 'variant_match': 1}
        unique_results.sort(key=lambda x: (x.get('score', 0), strategy_priority.get(x.get('search_strategy', ''), 0)), reverse=True)
        
        return unique_results[:top_k]
    
    def display_results(self, results, query):
        """显示搜索结果"""
        if not results:
            print("\n❌ 没有找到相关结果")
            return
        
        print(f"\n🎯 搜索结果 (共{len(results)}条):")
        print("=" * 80)
        
        for i, result in enumerate(results, 1):
            print(f"\n--- 结果 {i} ---")
            print(f"🏷️  产品: {result.get('product_code', 'N/A')}")
            print(f"📂 问题分类: {result.get('problem_category', 'N/A')}")
            print(f"🔍 搜索策略: {result.get('search_strategy', 'N/A')}")
            print(f"📊 匹配分数: {result.get('score', 0)}")
            print(f"🔄 变体类型: {result.get('variant_type', 'N/A')}")
            
            # 高亮显示匹配词
            original_text = result.get('original_text', '')
            search_terms = self.extract_search_terms(query)
            
            highlighted_text = original_text
            for term in search_terms:
                if len(term) > 1:
                    highlighted_text = highlighted_text.replace(term, f"【{term}】")
            
            print(f"📝 内容: {highlighted_text}")
            
            solution = result.get('solution', '')
            if solution:
                print(f"💡 解决方案: {solution}")
    
    def test_search_accuracy(self):
        """测试搜索准确性"""
        test_queries = [
            "PK511无产物",
            "模板质量问题",
            "PCR扩增失败",
            "投入量过高",
            "DNA降解",
            "建议重新制备",
            "500ng抑制反应"
        ]
        
        print("🧪 搜索准确性测试")
        print("=" * 50)
        
        for query in test_queries:
            print(f"\n🔍 测试查询: '{query}'")
            results = self.multi_strategy_search(query, top_k=5)
            print(f"✅ 找到 {len(results)} 条结果")
            
            if results:
                print(f"   最佳匹配: {results[0].get('problem_category', 'N/A')} - {results[0].get('search_strategy', 'N/A')}")


def interactive_enhanced_search():
    """交互式增强搜索"""
    print("🚀 增强版知识库搜索工具")
    print("=" * 50)
    print("💡 支持多策略搜索，提高命中率")
    print("=" * 50)
    
    searcher = EnhancedSearch()
    
    if not searcher.connect_to_milvus():
        return
    
    print("✅ 系统准备就绪！")
    print("💡 提示：输入 'test' 运行准确性测试，输入 'quit' 退出")
    
    while True:
        print("\n" + "=" * 50)
        query = input("🔍 请输入搜索问题: ").strip()
        
        if query.lower() in ['quit', 'exit', '退出']:
            print("👋 再见！")
            break
        
        if query.lower() == 'test':
            searcher.test_search_accuracy()
            continue
        
        if not query:
            print("⚠️  请输入有效的问题")
            continue
        
        try:
            results = searcher.multi_strategy_search(query, top_k=5)
            searcher.display_results(results, query)
        except Exception as e:
            print(f"❌ 搜索失败: {e}")


def command_line_enhanced_search():
    """命令行增强搜索"""
    if len(sys.argv) < 2:
        print("用法: python enhanced_search.py '搜索问题'")
        print("示例: python enhanced_search.py 'PK511无产物'")
        return
    
    query = ' '.join(sys.argv[1:])
    
    searcher = EnhancedSearch()
    if not searcher.connect_to_milvus():
        return
    
    results = searcher.multi_strategy_search(query, top_k=5)
    searcher.display_results(results, query)


if __name__ == "__main__":
    if len(sys.argv) > 1:
        command_line_enhanced_search()
    else:
        interactive_enhanced_search()
